from pathlib import Path
from wired_table_rec.utils.utils import VisTable
from table_cls import TableCls
from wired_table_rec.main import WiredTableInput, WiredTableRecognition
from lineless_table_rec.main import LinelessTableInput, LinelessTableRecognition

from rapidocr_paddle import RapidOCR

# 注意这里的参数
ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
# 初始化组件（只执行一次）
wired_input = WiredTableInput(device="cuda")
lineless_input = LinelessTableInput(device="cuda")
wired_engine = WiredTableRecognition(wired_input)
lineless_engine = LinelessTableRecognition(lineless_input)
table_cls = TableCls(model_type='paddle')
# from rapidocr import RapidOCR
# ocr_engine_cpu = RapidOCR()
viser = VisTable()
pipeline_yaml="src/config_yaml/OCR.yaml"
def infer_table_structure(model_type: str, img_path: str, save: bool = False):
    """
    执行表格识别流程，并可选保存识别结果

    Args:
        model_type (str): 使用模型 "wired" / "lineless" / "auto"
        img_path (str): 输入图片路径
        save (bool): 是否保存HTML与可视化图像

    Returns:
        table_results: 表格结构识别结果
    """
    # 自动选择模型
    if model_type == "auto":
        cls, _ = table_cls(img_path)
        table_engine = wired_engine if cls == "wired" else lineless_engine
    elif model_type == "wired":
        table_engine = wired_engine
    elif model_type == "lineless":
        table_engine = lineless_engine
    else:
        raise ValueError("model_type must be one of 'wired', 'lineless', 'auto'")

    # OCR提取
    # ocr_output = ocr_engine_cpu(img_path, return_word_box=True)
    # ocr_result_cpu = list(zip(ocr_output.boxes, ocr_output.txts, ocr_output.scores))

    ocr_result, _ = ocr_engine(img_path, return_word_box=True)
    import numpy as np
    ocr_result = [(np.array(box[0], dtype=np.float32),box[1],box[2]) for box in ocr_result]
    # 表格结构识别
    table_results = table_engine(img_path, ocr_result=ocr_result)

    # 保存
    if save:
        stem = Path(img_path).stem
        suffix = Path(img_path).suffix
        save_dir = Path("outputs")
        save_dir.mkdir(parents=True, exist_ok=True)

        save_html_path = save_dir / f"{stem}.html"
        save_drawed_path = save_dir / f"{stem}_table_vis{suffix}"
        save_logic_path = save_dir / f"{stem}_table_vis_logic{suffix}"

        viser(img_path, table_results, str(save_html_path), str(save_drawed_path), str(save_logic_path))

    return table_results

def batch_infer_table_structure(img_paths:list,model_type:str="auto",save:bool=False,device:str="gpu"):
    import sys
    import os

    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
    from paddlex_table.ocr import paddlex_ocr


    ocr_results = paddlex_ocr(img_paths, pipeline_yaml=pipeline_yaml, device=device)
    table_results = []
    for i, img_path in enumerate(img_paths):
        ocr_result = ocr_results[i]
        ocr_result = list(zip(ocr_result['rec_polys'],ocr_result['rec_texts'],ocr_result['rec_scores']))
        # 自动选择模型
        if model_type == "auto":
            cls, _ = table_cls(img_path)
            table_engine = wired_engine if cls == "wired" else lineless_engine
        elif model_type == "wired":
            table_engine = wired_engine
        elif model_type == "lineless":
            table_engine = lineless_engine
        else:
            raise ValueError("model_type must be one of 'wired', 'lineless', 'auto'")
        table_result = table_engine(img_path, ocr_result=ocr_result)
        table_results.append(table_result)
    return table_results

import time

if __name__ == "__main__":
    pipeline_yaml="../config_yaml/OCR.yaml"
    img = ["/home/fengjie/doc-parser/MinerU/src/TableStructureRec/40fb854e57b5258d90a73c4ad3ad52cf.png"]
    # img = ["/home/fengjie/doc-parser/MinerU/src/RapidTable/40fb854e57b5258d90a73c4ad3ad52cf.png","/home/fengjie/doc-parser/MinerU/src/RapidTable/3a417bf2cc6a33e21bf2cd995da6898bd1de20d6a35b24157b27a77244616b1d.jpg"]

    start_time = time.time()
    results = batch_infer_table_structure(model_type="auto", img_paths=img)
    end_time = time.time()
    
    elapsed = end_time - start_time
    print(f"识别完成，耗时：{elapsed:.3f} 秒")
    print(results)
    # print(results)
    

